from typing import Dict, Any, Optional, List
from dataclasses import dataclass, asdict
from datetime import datetime
from utils.logger import Logger
from utils.exceptions import AgentError, ToolError
from utils.constants import AGENT_STATUSES, TASK_STATUSES
from states.agent_state import AgentState
from states.task_state import TaskState
from .tools import text_qa_tools
from auxiliary_agents.compliance import check_compliance
from distributed.circuit_breaker import CircuitBreaker, circuit_breaker
from distributed.fallback import fallback_manager as fallback_handler

logger = Logger.get_logger(agent_name="text_qa_agent")


@dataclass
class TextQARequest:
    """文本问答请求数据结构"""
    query: str  # 问答查询文本
    user_id: str  # 用户ID
    task_id: str  # 任务ID
    metadata: Optional[Dict[str, Any]] = None  # 附加元数据（如上下文、优先级）


@dataclass
class TextQAResponse:
    """文本问答响应数据结构"""
    task_id: str
    answer: Optional[str]
    status: str  # 参考 TASK_STATUSES
    confidence: float  # 答案置信度（0-1）
    retrieved_sources: List[Dict[str, Any]]  # 检索来源（{text, similarity}）
    error_msg: Optional[str] = None
    timestamp: str = datetime.utcnow().isoformat() + "Z"


class TextQAAgent:
    """文本问答Agent：处理纯文本问答任务，集成检索增强能力"""

    def __init__(self):
        self.agent_id = "text_qa_agent"
        self.agent_state = AgentState(
            agent_type="text_qa",
            instance_id=self.agent_id
        )
        self.similarity_threshold = 0.6  # 相似度阈值（可通过配置中心动态调整）
        logger.info(f"TextQAAgent initialized: agent_id={self.agent_id}")

    def _update_agent_state(self, status: str):
        """更新Agent状态"""
        self.agent_state.status = status
        self.agent_state.last_heartbeat = datetime.utcnow()
        logger.debug(f"Agent state updated: {asdict(self.agent_state)}")

    def _validate_request(self, request: Dict[str, Any]) -> TextQARequest:
        """验证请求参数合法性"""
        required_fields = ["query", "user_id", "task_id"]
        for field in required_fields:
            if field not in request or not request[field]:
                raise AgentError(f"请求缺少必填字段: {field}")

        return TextQARequest(
            query=str(request["query"]).strip(),
            user_id=str(request["user_id"]).strip(),
            task_id=str(request["task_id"]).strip(),
            metadata=request.get("metadata", {})
        )

    @circuit_breaker(fallback=fallback_handler, threshold=5, timeout=30)
    def process_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """处理文本问答任务（核心入口）"""
        # 最明显的打印，确保能看到
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        print("!!!!!!!!! TextQAAgent.process_task 开始执行 !!!!!!!!!!!")
        print(f"!!!!!!!!! 任务数据: {task}")
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        
        task_id = task.get("task_id", "unknown_task")
        logger.info(f"开始处理文本问答任务: task_id={task_id}, query={task.get('query')}")

        try:
            self._update_agent_state(AGENT_STATUSES.BUSY)

            # 1. 验证请求
            request = self._validate_request(task)

            # 2. 合规性检查（过滤敏感查询）
            compliance_result = check_compliance(request.query)
            if not compliance_result["allowed"]:
                logger.warning(f"任务不合规: task_id={task_id}, reason={compliance_result['reason']}")
                response = TextQAResponse(
                    task_id=task_id,
                    answer=None,
                    status=TASK_STATUSES.REJECTED,
                    confidence=0.0,
                    retrieved_sources=[],
                    error_msg=f"查询不合规: {compliance_result['reason']}"
                )
                return asdict(response)

            # 3. 关键词提取（用于日志和后续优化）
            keywords = text_qa_tools.extract_keywords(request.query)
            logger.debug(f"提取关键词: task_id={task_id}, keywords={keywords}")

            # 4. 知识库检索
            print(f"!!!!!!!!! 开始调用 search_knowledge_base，查询: {request.query}")
            try:
                retrieved_texts = text_qa_tools.search_knowledge_base(request.query)
                print(f"!!!!!!!!! search_knowledge_base 调用成功，结果数量: {len(retrieved_texts)}")
                print(f"!!!!!!!!! 检索结果详情: {retrieved_texts}")
            except Exception as search_e:
                search_error_type = type(search_e).__name__
                search_error_msg = str(search_e)
                print(f"!!!!!!!!! search_knowledge_base 调用失败: {search_error_type}: {search_error_msg}")
                import traceback
                print(f"!!!!!!!!! 检索错误堆栈:\n{traceback.format_exc()}")
                raise
            logger.debug(f"检索结果: task_id={task_id}, count={len(retrieved_texts)}")

            # 5. 生成答案
            print(f"!!!!!!!!! 开始调用 generate_answer，相似度阈值: {self.similarity_threshold}")
            try:
                answer = text_qa_tools.generate_answer(
                    query=request.query,
                    retrieved_texts=retrieved_texts,
                    similarity_threshold=self.similarity_threshold
                )
                print(f"!!!!!!!!! generate_answer 调用成功，生成答案: {answer}")
            except Exception as gen_e:
                gen_error_type = type(gen_e).__name__
                gen_error_msg = str(gen_e)
                print(f"!!!!!!!!! generate_answer 调用失败: {gen_error_type}: {gen_error_msg}")
                import traceback
                print(f"!!!!!!!!! 生成错误堆栈:\n{traceback.format_exc()}")
                raise

            # 6. 计算置信度（基于检索结果的最高相似度）
            confidence = max([score for _, score in retrieved_texts], default=0.0) if answer else 0.0

            # 7. 构造响应
            response = TextQAResponse(
                task_id=task_id,
                answer=answer or "未找到相关信息，无法回答",
                status=TASK_STATUSES.SUCCESS if answer else TASK_STATUSES.FAILED,
                confidence=confidence,
                retrieved_sources=[
                    {"text": text, "similarity": score}
                    for text, score in retrieved_texts
                ]
            )

            logger.info(f"任务处理完成: task_id={task_id}, status={response.status}, confidence={confidence}")
            return asdict(response)

        except ToolError as e:
            print(f"!!!!!!!!! 捕获到 ToolError: {str(e)}")
            logger.error(f"工具调用失败: task_id={task_id}, error={str(e)}")
            response = TextQAResponse(
                task_id=task_id,
                answer=None,
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                retrieved_sources=[],
                error_msg=f"工具调用失败: {str(e)}"
            )
            return asdict(response)

        except AgentError as e:
            print(f"!!!!!!!!! 捕获到 AgentError: {str(e)}")
            logger.error(f"Agent处理失败: task_id={task_id}, error={str(e)}")
            response = TextQAResponse(
                task_id=task_id,
                answer=None,
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                retrieved_sources=[],
                error_msg=f"任务处理失败: {str(e)}"
            )
            return asdict(response)

        except Exception as e:
            error_type = type(e).__name__
            print(f"!!!!!!!!! 捕获到 未知错误: {error_type}: {str(e)}")
            import traceback
            print(f"!!!!!!!!! 未知错误堆栈:\n{traceback.format_exc()}")
            logger.error(f"未知错误: task_id={task_id}, error={str(e)}", exc_info=True)
            response = TextQAResponse(
                task_id=task_id,
                answer=None,
                status=TASK_STATUSES.FAILED,
                confidence=0.0,
                retrieved_sources=[],
                error_msg=f"系统错误: {str(e)}"
            )
            return asdict(response)

        finally:
            print(f"!!!!!!!!! 执行 finally 块，更新状态为 RUNNING")
            self._update_agent_state(AGENT_STATUSES.RUNNING)

    def get_agent_state(self) -> Dict[str, Any]:
        """获取Agent当前状态（供监控和调度中心调用）"""
        return asdict(self.agent_state)


# 单例Agent实例（全局复用）
text_qa_agent = TextQAAgent()

# 测试入口（本地调试用）
if __name__ == "__main__":
    test_task = {
        "task_id": "test_001",
        "query": "什么是分布式多Agent系统？",
        "user_id": "test_user",
        "metadata": {"priority": "normal"}
    }

    result = text_qa_agent.process_task(test_task)
    print("测试结果:", result)